# ----------------------------------------------------------------------------------
# Microsoft Developer & Platform Evangelism
# 
# Copyright (c) Microsoft Corporation. All rights reserved.
# 
# THIS CODE AND INFORMATION ARE PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND, 
# EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE IMPLIED WARRANTIES 
# OF MERCHANTABILITY AND/OR FITNESS FOR A PARTICULAR PURPOSE.
# ----------------------------------------------------------------------------------
# The example companies, organizations, products, domain names,
# e-mail addresses, logos, people, places, and events depicted
# herein are fictitious.  No association with any real company,
# organization, product, domain name, email address, logo, person,
# places, or events is intended or should be inferred.
# ----------------------------------------------------------------------------------
#
# THIS IS SAMPLE CODE. BUGS ARE EXPECTED. SECURITY IS NOT.
#
# This module provides autoscaling functionality 
#

from datetime import *
from time import *
from subprocess import *
from ast import *
from azure.storage import *
from socket import *
from azure import *
from azure.servicemanagement import *
from os import *
from sys import *
from shlex import *

#PARAMETERS
#name of queue in
qname = 'inqueue'
#subscription and certificate
subscription_id = 'enter your subscription id'
certificate_path ="CURRENT_USER\\my\\ enter your certificate name"
#Name of service to create
name = "geres"
#Location of service
location = "West US"
#storage account
storage_account = "enter your storage account"
storage_key = 'your storage key'
#vhd image to use, as returned by list_os_images
image_name = " your vhd image, duly sysprepped and containing python listener "
#max number of instances to use. Choose appopriate n. for your subscription
maxinstances = 3
#size of instances
vmsize = "Small"
#tolerable waiting time for jobs in queue
tolerable = timedelta(minutes=1)

# MAIN PROGRAM
# instantiate a queue service
queue_service = QueueService(account_name=storage_account, account_key=storage_key)
# instantiate a service management service
sms = ServiceManagementService(subscription_id, certificate_path)


# NOTE: the following functions do not work; there is a bug somewhere in the sdk where add_role is defined.
# Endpoints are not created, networks not joined. The function is left here as an example of how things should work.
# The functions are replaced by a Powershell script, which does work but makes this code impossible to port to Linux. As things are fixed in the SDK
# we'll fix this code too.

#function to build several VMs at once
def buildvms(name, number, location, size, image):
    # Windows VM configuration
    windows_config = WindowsConfigurationSet(name, 'Pacman10')
    windows_config.domain_join = None
    #linux_config = LinuxConfigurationSet(name, 'topguy', 'Pacman10', True)

    endpoint_config = ConfigurationSet()
    endpoint_config.configuration_set_type='NetworkConfiguration'

    endpoint1=ConfigurationSetInputEndpoint(name = 'rdp', protocol = 'tcp', port = '33890', local_port = '3389', load_balanced_endpoint_set_name = None , enable_direct_server_return = False)
    endpoint2=ConfigurationSetInputEndpoint(name = 'web', protocol = 'tcp', port = '8080', local_port = '80', load_balanced_endpoint_set_name = None , enable_direct_server_return = False)

    endpoint_config.input_endpoints.input_endpoints.append(endpoint1)
    endpoint_config.input_endpoints.input_endpoints.append(endpoint2)

    #build the services
    index=1
    while index <= number:
        servicename= name+str(index)
        sms.create_hosted_service(service_name=servicename, label=servicename,location=location)
    # Destination storage account container/blob where the VM disk will be created
        media_link = 'http://gmstorage.blob.core.windows.net/vhds/'+servicename+'.vhd'
        os_hd = OSVirtualHardDisk(image, media_link)
     
        sms.create_virtual_machine_deployment(service_name=servicename,
        deployment_name=servicename,
        deployment_slot='production',
        label=servicename,
        role_name=servicename,
        system_config=windows_config,
        network_config = endpoint_config,
        os_virtual_hard_disk=os_hd,
        role_size=size)
        result = sms.get_deployment_by_name(service_name=servicename, deployment_name = servicename)
        print 'Deployment: ', servicename, ' Status: ', result.status 
        index = index +1

def buildvms2(name, number, location, size, image):
    # Windows VM configuration
    windows_config = WindowsConfigurationSet(name, 'Pacman10')
    windows_config.domain_join = None
    #linux_config = LinuxConfigurationSet(name, 'topguy', 'Pacman10', True)

    endpoint_config = ConfigurationSet()
    endpoint_config.configuration_set_type='NetworkConfiguration'

    endpoint1=ConfigurationSetInputEndpoint(name = 'rdp', protocol = 'tcp', port = '33890', local_port = '3389', load_balanced_endpoint_set_name = None , enable_direct_server_return = False)
    endpoint2=ConfigurationSetInputEndpoint(name = 'web', protocol = 'tcp', port = '8080', local_port = '80', load_balanced_endpoint_set_name = None , enable_direct_server_return = False)

    endpoint_config.input_endpoints.input_endpoints.append(endpoint1)
    endpoint_config.input_endpoints.input_endpoints.append(endpoint2)

    sms.create_hosted_service(service_name=name, label=name,location=location)
    
    media_link = 'http://gmstorage.blob.core.windows.net/vhds/'+name+'.vhd'
    os_hd = OSVirtualHardDisk(image, media_link)
    
    servicename = name 
    #deploy the first machine of the role, which creates deployment
    sms.create_virtual_machine_deployment(service_name=servicename,
        deployment_name=servicename,
        deployment_slot='production',
        label=name,
        role_name=name+'1',
        system_config=windows_config,
        network_config = endpoint_config,
        os_virtual_hard_disk=os_hd,
        role_size=size)
        
    result = sms.get_deployment_by_name(service_name=servicename, deployment_name = name)
   #pause required to avoid azure api bugs that cause conflict errors
    while result.status != 'Running':
        print 'Deployment: ', name, ' Status: ', result.status 
        sleep (60)
        result = sms.get_deployment_by_name(service_name=servicename, deployment_name = name)
    #build the remaining VMs
    index=2
    while index <= number:
        namei = servicename + str(index)
        media_link = 'http://gmstorage.blob.core.windows.net/vhds/'+namei+'.vhd'
        os_hd = OSVirtualHardDisk(image, media_link)
        sms.add_role(service_name = servicename, deployment_name = servicename, role_name = namei, system_config=windows_config, os_virtual_hard_disk=os_hd, role_size=size)
        #again pause required because of generally slow azure management api
        result = sms.get_deployment_by_name(service_name=servicename, deployment_name = servicename)
        print 'Deployment: ', servicename, ' Status: ', result.status 
        sleep (60)
        index = index +1

# since VMs are free until turned on as of June 2013, create as many VMs as requested by maxcores, but leave them off until needed.
try:
    result=sms.get_hosted_service_properties(service_name = name)
    #not really useful to list everything, but it helps me understand what is running.
    #result=sms.list_hosted_services()
    #for item in result:
    print 'Service name:', result.service_name, 'Status: ', result.hosted_service_properties.status   
except WindowsAzureMissingResourceError as error:
    print 'Service ', name, ' does not exist. Creating...'
    #buildvms2(name, maxcores, location, 'Small', image_name)
    command = "C:/WINDOWS/system32/WindowsPowerShell/v1.0/powershell.exe -ExecutionPolicy Unrestricted ./createvms.ps1 "+name+" "+str(maxinstances)+" "+location+" "+storage_account+" "+vmsize+" "+image_name
    args = split(command)
    #note how the python SDK splits arguments with a space into two different arguments, e.g. "West US" becomes "West" "US".
    #this must be fixed in the receiving powershell script. There is no python workaround.
    result = call(args)
    if (int(result) == 0):
        print "Oops.. No VM deployed! Something went wrong. Exiting..."
        exit(1)

nofmsg=0
try:
    qmetadata=queue_service.get_queue_metadata(qname)
    nofmsg = int(qmetadata['x-ms-approximate-messages-count'])
    print ("Messages in queue: ", nofmsg)
except WindowsAzureMissingResourceError as error:
    print 'Queue ', qname, 'does not exist. The program cannot continue.'
    exit(1)


while True:
    #get queue properties and n. of messages
    qmetadata = queue_service.get_queue_metadata(qname)
    nofmsg = qmetadata['x-ms-approximate-messages-count']
    print ("Messages in queue: ", nofmsg)

    # This is the part that decides to start VMs or not
    while int(nofmsg) <= 0:
        # wait 10 sec
        sleep(10)
        qmetadata = queue_service.get_queue_metadata(qname)
        nofmsg = qmetadata['x-ms-approximate-messages-count']
        print 'Messages in queue: ', nofmsg

    #if there is a running VM (or more), try and understand how long the job has been waiting
    waittime = timedelta(minutes=0)
    print 'Tolerable wait time is: ', tolerable, ' minutes.'

    #if the waittime is tolerable, do nothing
    while waittime <= tolerable:
        messages = queue_service.peek_messages(qname)
        for message in messages:
            firstjob = literal_eval(message.message_text)
            waittime= datetime.now() -  datetime.strptime(firstjob['TimeIn'], '%Y-%m-%dT%H:%M:%S.%f')
            print ('First job wait time: {0} seconds'.format(abs(waittime.total_seconds())))
        sleep(10)

    #if the waittime is not tolerable, look for VMs in defined services that are not running and start them
    #note that the vms will shut themselves down if there is no work, so there is no need to keep track of idle or active ones.
    #NOTE: there is a bug in the February2013 implementation of the Python SDK, where the StoppedDeallocated status is NOT recognized and such vm disappears off the list
    #only way around is to use Powershell June2013 or newer
    #get-azurevm -ServiceName <service name> | where {$_.instancestatus -in ('Stopped', 'StoppedDeallocated')}
    
        #this is how it should work once they fix the python SDK
        #properties = sms.get_hosted_service_properties(service_name=name, embed_detail = True)
        #for deployment in properties.deployments:
        #        print '\tDeployment: ', deployment.name, '\tStatus ', deployment.status
        #        for role in deployment.role_instance_list:
        #            print '\t\tRole : ', role.role_name, '\tInstance: ', role.instance_name, '\tStatus: ', role.instance_status
        #            #if there are roles in the service named before that are not running or about to run, start them
        #            if (role.instance_status in ('StoppedVM', 'StoppedDeallocated')) and (startedone==False):
        #                    print '\t\tStarting ', role.role_name, 'Instance ', role.instance_name
        #                    sms.start_role(vm.service_name, deployment.name, role.role_name)
        #                    #start one role at a time, then check for latency again
        #                    startedone=True
        
    #This is the Powershell workaround
    #The script will start 1 VM at a time. This is done to minimize additional cost; you can start as many as you need.
    result = call (["C:\\WINDOWS\\system32\\WindowsPowerShell\\v1.0\\powershell.exe", "-ExecutionPolicy", "Unrestricted", "./returnstoppedvmsn.ps1", name], cwd=getcwd())
    print 'Stopped VMs = ', int(result)
    if (int(result)>0):
        print 'Started 1 VM to pick up the slack'
    else:
        print 'There are no VMs left to start. Either increase n. of cores in subscription'
        print 'or wait for tasks to finish.'
        sleep(10)
          
                            


                
        


